package com.tianji.aigc.memory;

import cn.hutool.core.collection.CollStreamUtil;
import cn.hutool.json.JSONNull;
import cn.hutool.json.JSONUtil;
import com.tianji.common.utils.CollUtils;
import com.tianji.common.utils.JsonUtils;
import com.tianji.common.utils.UserContext;
import jakarta.annotation.Resource;
import lombok.RequiredArgsConstructor;
import org.springframework.ai.chat.memory.ChatMemory;
import org.springframework.ai.chat.messages.Message;
import org.springframework.beans.factory.annotation.Autowired;
import org.springframework.data.redis.core.BoundListOperations;
import org.springframework.data.redis.core.StringRedisTemplate;

import java.util.List;

/**
 * @Description: 使用Redis作为存储会话记忆的容器
 * @Author: lemon
 * @CreateTime: 2025-06-09  20:53
 * @Version: 1.0
 */

public class RedisChatMemory implements ChatMemory {

    @Autowired
    private StringRedisTemplate redisTemplate;

    private static final String DEFAULT_PREFIX = "CHAT:";

    private final String prefix;

    public RedisChatMemory() {
        this.prefix = DEFAULT_PREFIX;
    }

    public RedisChatMemory(String prefix) {
        this.prefix = prefix;
    }

    @Override
    public void add(String conversationId, List<Message> messages) {
        if(CollUtils.isEmpty(messages)){
            return;
        }
        String key = getMemoryKey(conversationId);
        BoundListOperations<String, String> listOps = redisTemplate.boundListOps(key);
        messages.forEach(message -> listOps.rightPush(MessageUtil.toJson(message)));
    }

    /**
     * 根据会话ID获取指定数量的最新消息
     *
     * @param conversationId 会话唯一标识符
     * @param lastN          需要获取的最后N条消息数量（N>0）
     * @return 包含消息对象的列表，若lastN<=0则返回空列表
     */
    @Override
    public List<Message> get(String conversationId, int lastN) {
        // 验证参数有效性，当lastN非正数时直接返回空结果
        if (lastN <= 0) {
            return List.of();
        }
        String key = getMemoryKey(conversationId);
        List<String> messages = redisTemplate.opsForList().range(key, 0, lastN);
        return CollStreamUtil.toList(messages, MessageUtil::toMessage);
    }

    @Override
    public void clear(String conversationId) {
        redisTemplate.delete(getMemoryKey(conversationId));
    }

    public String getMemoryKey(String conversationId) {
        return prefix + conversationId;
    }
}
